# =========================================================================
# === Road Network Repair & Visualization Script (v6 - Optimized) =========
# =========================================================================
# 描述:
#   这是一个性能和结构都经过优化的版本。核心改动：
#   1. [性能] 不再在循环中反复重建网络图，而是在内存中动态更新。
#   2. [结构] 将修复逻辑封装在独立的函数中，主流程更清晰。
#   3. [健壮性] 增强了对输入数据的检查。
# =========================================================================

import geopandas as gpd
from shapely.geometry import Point, LineString
import networkx as nx
from scipy.spatial import cKDTree
import pandas as pd
import warnings
import time
import folium
import sys
import traceback

# --- 0. 环境设置 ---
warnings.simplefilter(action='ignore', category=FutureWarning)

# =========================================================================
# === 1. 参数定义 (保持不变) ==============================================
# =========================================================================
INPUT_ROAD_SHAPEFILE = "2025年福建省道路数据(修复).shp" 
INPUT_BOUNDARY_SHAPEFILE = "仓山区最终完整边界.shp"
OUTPUT_REPAIRED_SHAPEFILE = "仓山区修复后路网.shp"
OUTPUT_VISUALIZATION_MAP = "repair_visualization_map.html"
CLIP_BUFFER_DISTANCE = 3000
PROJECTION_EPSG = 3857
REPAIR_THRESHOLDS = [50, 150, 300]

# =========================================================================
# === 2. 辅助函数 =========================================================
# =========================================================================
def gdf_to_nx(gdf: gpd.GeoDataFrame) -> nx.Graph:
    """从GeoDataFrame高效构建NetworkX图"""
    G = nx.Graph()
    for geom in gdf.geometry:
        if geom is None or not isinstance(geom, LineString):
            continue
        for i in range(len(geom.coords) - 1):
            p1, p2 = geom.coords[i], geom.coords[i+1]
            G.add_edge(p1, p2)
    return G

def repair_network_iteratively(G: nx.Graph, thresholds: list) -> list:
    """
    在NetworkX图上执行多轮拓扑修复，并返回所有新建的“桥梁”。
    """
    all_new_bridges = []
    
    for i, threshold in enumerate(thresholds):
        print(f"  --- [第 {i+1}/{len(thresholds)} 轮修复] 使用阈值: {threshold} 米 ---")
        
        if nx.is_connected(G):
            print("  - ✅ 路网在本轮已完全连通，提前停止修复。")
            break
            
        components = sorted(list(nx.connected_components(G)), key=len, reverse=True)
        mainland_nodes = components[0]
        island_nodes = set().union(*(c for c in components[1:]))
        
        if not island_nodes:
            print("  - 未发现更多“岛屿”，修复完成。")
            break

        print(f"    - 正在为 {len(island_nodes)} 个孤立节点查找连接点...")
        mainland_kdtree = cKDTree(list(mainland_nodes))
        island_node_list = list(island_nodes)
        
        distances, mainland_indices = mainland_kdtree.query(island_node_list, k=1)
        
        round_bridges = []
        for j, dist in enumerate(distances):
            if dist <= threshold and dist > 0:
                island_node = island_node_list[j]
                mainland_node = list(mainland_nodes)[mainland_indices[j]]
                
                # 直接在图中添加边（建桥），而不是创建几何对象
                G.add_edge(island_node, mainland_node)
                # 记录新建的桥梁几何，用于后续可视化和保存
                round_bridges.append(LineString([island_node, mainland_node]))
        
        if round_bridges:
            print(f"    - 本轮成功构建 {len(round_bridges)} 座新“桥梁”。")
            all_new_bridges.extend(round_bridges)
        else:
            print("    - 本轮未发现可在阈值内连接的断点。")
            
    return all_new_bridges

# =========================================================================
# === 3. 主执行流程 =======================================================
# =========================================================================
def main():
    start_time = time.time()
    print("--- 开始执行路网裁剪与【多轮】拓扑修复 (优化版) ---")

    try:
        # --- 步骤 1: 加载并预处理数据 ---
        print(f"\n[步骤 1/4] 正在加载并预处理数据...")
        
        road_gdf_raw = gpd.read_file(INPUT_ROAD_SHAPEFILE)
        boundary_gdf_raw = gpd.read_file(INPUT_BOUNDARY_SHAPEFILE)
        if boundary_gdf_raw.crs is None:
            print("  - ⚠️ 警告：边界文件缺少CRS，将强制假定其为WGS-84 (EPSG:4326)。")
            boundary_gdf_raw.set_crs("EPSG:4326", inplace=True)
        
        road_gdf_proj = road_gdf_raw.to_crs(f"EPSG:{PROJECTION_EPSG}")
        boundary_gdf_proj = boundary_gdf_raw.to_crs(f"EPSG:{PROJECTION_EPSG}")
        
        print("  - ✅ 原始数据加载成功。")
        
        # --- 步骤 2: 裁剪与几何清理 ---
        print(f"\n[步骤 2/4] 正在裁剪并清理路网几何...")
        boundary_polygon = boundary_gdf_proj.unary_union
        clip_geometry = boundary_polygon.buffer(CLIP_BUFFER_DISTANCE) if CLIP_BUFFER_DISTANCE > 0 else boundary_polygon
        
        road_gdf_clipped = gpd.clip(road_gdf_proj, clip_geometry)
        road_gdf_simplified = road_gdf_clipped.copy()
        road_gdf_simplified['geometry'] = road_gdf_clipped.geometry.simplify(tolerance=10)
        road_gdf_simplified = road_gdf_simplified[~road_gdf_simplified.geometry.is_empty]
        road_gdf_exploded = road_gdf_simplified.explode(index_parts=False)
        print(f"  - ✅ 路网裁剪和清理完成，得到 {len(road_gdf_exploded)} 条路段。")

        # --- 步骤 3: 构建图并执行拓扑修复 ---
        print(f"\n[步骤 3/4] 正在构建初始网络图并执行修复...")
        
        # a. 只构建一次初始图
        initial_graph = gdf_to_nx(road_gdf_exploded)
        
        # b. 在图上执行多轮修复
        new_bridges_geoms = repair_network_iteratively(initial_graph, REPAIR_THRESHOLDS)
        
        # c. 将新建的桥梁合并回GeoDataFrame
        if new_bridges_geoms:
            bridges_gdf = gpd.GeoDataFrame(geometry=new_bridges_geoms, crs=road_gdf_exploded.crs)
            road_gdf_repaired = pd.concat([road_gdf_exploded, bridges_gdf], ignore_index=True)
        else:
            road_gdf_repaired = road_gdf_exploded
        
        # d. 保存修复后的文件
        road_gdf_repaired.to_file(OUTPUT_REPAIRED_SHAPEFILE, driver='ESRI Shapefile', encoding='utf-8')
        print(f"\n✅ [成功] 已将最终修复好的路网 ({len(road_gdf_repaired)}条) 保存到: '{OUTPUT_REPAIRED_SHAPEFILE}'")

        # --- 步骤 4: 生成可视化地图 ---
        print(f"\n[步骤 4/4] 正在生成可视化修复效果的地图...")
        
        boundary_map_gdf = boundary_gdf_proj.to_crs("EPSG:4326")
        original_roads_map_gdf = road_gdf_exploded.to_crs("EPSG:4326")
        map_center = [boundary_map_gdf.unary_union.centroid.y, boundary_map_gdf.unary_union.centroid.x]
        
        m = folium.Map(location=map_center, zoom_start=12, tiles="OpenStreetMap")
        
        folium.GeoJson(original_roads_map_gdf, name="修复前的路网", style_function=lambda x: {'color': 'gray', 'weight': 1.5}).add_to(m)
        
        if new_bridges_geoms:
            bridges_map_gdf = gpd.GeoDataFrame(geometry=new_bridges_geoms, crs=road_gdf_exploded.crs).to_crs("EPSG:4326")
            folium.GeoJson(bridges_map_gdf, name=f"新建的连接“桥梁” (共 {len(new_bridges_geoms)} 座)", style_function=lambda x: {'color': 'red', 'weight': 3, 'dashArray': '5, 5'}).add_to(m)
        
        folium.GeoJson(boundary_map_gdf, name="研究区域边界", style_function=lambda x: {'color': 'black', 'weight': 2, 'fillOpacity': 0}).add_to(m)
        
        bounds = boundary_map_gdf.total_bounds
        m.fit_bounds([[bounds[1], bounds[0]], [bounds[3], bounds[2]]])
        folium.LayerControl().add_to(m)
        m.save(OUTPUT_VISUALIZATION_MAP)
        print(f"  - ✅ [成功] 可视化地图已保存到: '{OUTPUT_VISUALIZATION_MAP}'")

        total_time = time.time() - start_time
        print(f"\n--- 全部流程完成，总耗时: {total_time:.2f} 秒 ---")

    except Exception as e:
        print(f"\n❌ [严重错误] 在主流程中发生错误: {e}")
        traceback.print_exc()
        sys.exit(1)

if __name__ == "__main__":
    main()